// ComHook.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include "HookUtilities.h"

#include <stdio.h>
#include <stdarg.h>
#include <time.h>

#define HELENA_CUSTOMIMPORT_REGISTERY_KEY "SOFTWARE\\Helena\\Custom Importer"

typedef HFILE (WINAPI *LPOpenFile)(LPCSTR lpFileName, LPOFSTRUCT lpReOpenBuff, UINT uStyle);
typedef HANDLE (WINAPI *LPCreateFileA)(LPCSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile);
typedef HANDLE (WINAPI *LPCreateFileW)(LPCWSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile);
typedef BOOL (WINAPI *LPCloseHandle)(HANDLE hObject);
typedef BOOL (WINAPI *LPWriteFile)(HANDLE hFile, LPCVOID lpBuffer, DWORD nNumberOfBytesToWrite, LPDWORD lpNumberOfBytesWritten, LPOVERLAPPED lpOverlapped);
typedef BOOL (WINAPI *LPWriteFileEx)(HANDLE hFile, LPCVOID lpBuffer, DWORD nNumberOfBytesToWrite, LPOVERLAPPED lpOverlapped, LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
typedef BOOL (WINAPI *LPReadFile)(HANDLE hFile, LPVOID lpBuffer, DWORD nNumberOfBytesToRead, LPDWORD lpNumberOfBytesRead, LPOVERLAPPED lpOverlapped);
typedef BOOL (WINAPI *LPReadFileEx)(HANDLE hFile, LPVOID lpBuffer, DWORD nNumberOfBytesToRead, LPOVERLAPPED lpOverlapped, LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);

LPOpenFile g_lpOpenFile = NULL;
LPCreateFileA g_lpCreateFileA = NULL;
LPCreateFileW g_lpCreateFileW = NULL;
LPCloseHandle g_lpCloseHandle = NULL;
LPWriteFile g_lpWriteFile = NULL;
LPWriteFileEx g_lpWriteFileEx = NULL;
LPReadFile g_lpReadFile = NULL;
LPReadFileEx g_lpReadFileEx = NULL;

HANDLE g_hLogFile = INVALID_HANDLE_VALUE;

BOOL Log(LPCSTR message, ...);

HFILE WINAPI OpenFile_Hook(LPCSTR lpFileName, LPOFSTRUCT lpReOpenBuff, UINT uStyle)
{
	Log("Open File %s", lpFileName);
	return g_lpOpenFile(lpFileName, lpReOpenBuff, uStyle);
}

HANDLE WINAPI CreateFileA_Hook(LPCSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	HANDLE handle = g_lpCreateFileA(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);
	Log("Create File (0x%x): %s", handle, lpFileName);
	return handle;
}

HANDLE WINAPI CreateFileW_Hook(LPCWSTR lpFileName, DWORD dwDesiredAccess, DWORD dwShareMode, LPSECURITY_ATTRIBUTES lpSecurityAttributes, DWORD dwCreationDisposition, DWORD dwFlagsAndAttributes, HANDLE hTemplateFile)
{
	HANDLE handle = g_lpCreateFileW(lpFileName, dwDesiredAccess, dwShareMode, lpSecurityAttributes, dwCreationDisposition, dwFlagsAndAttributes, hTemplateFile);
	Log("Create File (0x%x): %s", handle, lpFileName);
	return handle;
}

BOOL WINAPI CloseHandle_Hook(HANDLE hObject)
{
	Log("Close Handle (0x%x)", hObject);
	return g_lpCloseHandle(hObject);
}

BOOL WINAPI WriteFile_Hook(HANDLE hFile, LPCVOID lpBuffer, DWORD nNumberOfBytesToWrite, LPDWORD lpNumberOfBytesWritten, LPOVERLAPPED lpOverlapped)
{
	char* pBuffer = new char[nNumberOfBytesToWrite + 1];
	memcpy(pBuffer, lpBuffer, nNumberOfBytesToWrite);
	pBuffer[nNumberOfBytesToWrite] = 0;
	Log("Write File (0x%x):\r\n%s", hFile, pBuffer);
	delete[] pBuffer;

	return g_lpWriteFile(hFile, lpBuffer, nNumberOfBytesToWrite, lpNumberOfBytesWritten, lpOverlapped);
}

BOOL WINAPI WriteFileEx_Hook(HANDLE hFile, LPCVOID lpBuffer, DWORD nNumberOfBytesToWrite, LPOVERLAPPED lpOverlapped, LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	char* pBuffer = new char[nNumberOfBytesToWrite + 1];
	memcpy(pBuffer, lpBuffer, nNumberOfBytesToWrite);
	pBuffer[nNumberOfBytesToWrite] = 0;
	Log("Write File EX (0x%x):\r\n%s", hFile, pBuffer);
	delete[] pBuffer;

	return g_lpWriteFileEx(hFile, lpBuffer, nNumberOfBytesToWrite, lpOverlapped, lpCompletionRoutine);
}

BOOL WINAPI ReadFile_Hook(HANDLE hFile, LPVOID lpBuffer, DWORD nNumberOfBytesToRead, LPDWORD lpNumberOfBytesRead, LPOVERLAPPED lpOverlapped)
{
	BOOL bResult = g_lpReadFile(hFile, lpBuffer, nNumberOfBytesToRead, lpNumberOfBytesRead, lpOverlapped);

	char* pBuffer = new char[*lpNumberOfBytesRead + 1];
	memcpy(pBuffer, lpBuffer, *lpNumberOfBytesRead);
	pBuffer[*lpNumberOfBytesRead] = 0;
	Log("Read File (0x%x):\r\n%s", hFile, pBuffer);
	delete[] pBuffer;

	return bResult;
}

BOOL WINAPI ReadFileEx_Hook(HANDLE hFile, LPVOID lpBuffer, DWORD nNumberOfBytesToRead, LPOVERLAPPED lpOverlapped, LPOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine)
{
	BOOL bResult = g_lpReadFileEx(hFile, lpBuffer, nNumberOfBytesToRead, lpOverlapped, lpCompletionRoutine);
	Log("Read File EX (0x%x)", hFile);
	return bResult;
}

BOOL g_bLoggingEnabled;
BOOL g_bLogDetails;
char g_strInstallationDirectory[1024];
char g_strDatabaseConnectionString[1024];

BOOL SetInitialisationFatalError(LPCSTR strMessage)
{
	HKEY hKey;
    if (ERROR_SUCCESS == RegOpenKeyExA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, 0, KEY_ALL_ACCESS , &hKey))
	{
		RegSetValueExA(hKey, "Last Fatal Error", 0, REG_SZ, (LPBYTE)strMessage, strlen(strMessage));
		RegCloseKey(hKey);
	}

	return FALSE;
}

BOOL CreateLogFile(LPCSTR strLoggingDirectory)
{
	try
	{
		char strFilename[1024];
		strcpy(strFilename, strLoggingDirectory);
		int iTail = strlen(strFilename);
		if (iTail > 1 && strFilename[iTail - 1] != '\\') { strFilename[iTail++] = '\\'; strFilename[iTail] = '\0'; }

		time_t now = time(NULL);
		tm tmTime = *(localtime(&now));
		strftime(strFilename + iTail, 20, "%Y%m%d_%H%M%S.LOG", &tmTime);

		g_hLogFile = CreateFileA(strFilename, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, INVALID_HANDLE_VALUE);

		if (g_hLogFile == INVALID_HANDLE_VALUE)
			return SetInitialisationFatalError("Failed creating logging file");
	}
	catch (...)
	{
		return SetInitialisationFatalError("Unhandled exception while creating log file");
	}
}

BOOL Log(LPCSTR message, ...)
{
	try
	{
		if (g_bLoggingEnabled && g_hLogFile != INVALID_HANDLE_VALUE)
		{
			va_list pArg;
			va_start(pArg, message);
			char buf[2048];
			_vsnprintf(buf, 2048, message, pArg);
			va_end(pArg);

			DWORD dwBytesWritten;
			if (FALSE == g_lpWriteFile(g_hLogFile, buf, strlen(buf), &dwBytesWritten, NULL))
				return SetInitialisationFatalError("Unhandled Exception while logging");
			if (FALSE == g_lpWriteFile(g_hLogFile, "\r\n", 2, &dwBytesWritten, NULL))
				return SetInitialisationFatalError("Unhandled Exception while logging");
			if (FALSE == FlushFileBuffers(g_hLogFile))
				return SetInitialisationFatalError("Unhandled Exception while logging");

			return TRUE;
		}
	}
	catch (...)
	{
		return SetInitialisationFatalError("Unhandled Exception while logging");
	}
}

LPCSTR GetProcessName()
{
	char filename[1024];
	GetModuleFileNameA(GetModuleHandle(NULL), filename, sizeof(filename));
	strupr(filename);

	char* filenameStart;

	for (filenameStart = filename; *filenameStart != 0; filenameStart++) ;
	for (; *filenameStart != '\\'; filenameStart--) ;
	return filenameStart + 1;
}

BOOL ReadConfiguration()
{
	char strText[1024]; BYTE bytBoolean;

	DWORD dwBufferSize = sizeof(bytBoolean);
	if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Use Custom Import", RRF_RT_ANY, NULL, &bytBoolean, &dwBufferSize))
		return FALSE;
	if (bytBoolean == 0)
		return FALSE;

	dwBufferSize = sizeof(strText);
	if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Hook Target", RRF_RT_ANY, NULL, strText, &dwBufferSize))
		return FALSE;
	strupr(strText);
	if (strcmp(GetProcessName(), strText) != 0)
		return FALSE;

	dwBufferSize = sizeof(g_strInstallationDirectory);
	if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Installation Directory", RRF_RT_ANY, NULL, g_strInstallationDirectory, &dwBufferSize))
		return FALSE;

	dwBufferSize = sizeof(bytBoolean);
	if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Logging Enabled", RRF_RT_ANY, NULL, &bytBoolean, &dwBufferSize))
		return FALSE;
	g_bLoggingEnabled = bytBoolean != 0;

	if (g_bLoggingEnabled)
	{
		dwBufferSize = sizeof(strText);
		if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Logging Directory", RRF_RT_ANY, NULL, strText, &dwBufferSize))
			return FALSE;
		if (!CreateLogFile(strText))
			return FALSE;

		dwBufferSize = sizeof(bytBoolean);
		if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Log Details", RRF_RT_ANY, NULL, &bytBoolean, &dwBufferSize))
			return FALSE;
		g_bLogDetails = bytBoolean != 0;
	}

	dwBufferSize = sizeof(g_strDatabaseConnectionString);
	if (ERROR_SUCCESS != RegGetValueA(HKEY_CURRENT_USER, HELENA_CUSTOMIMPORT_REGISTERY_KEY, "Database Information", RRF_RT_ANY, NULL, g_strDatabaseConnectionString, &dwBufferSize))
		return FALSE;

	return TRUE;
}

void InitialiseHooks()
{
	if (!ReadConfiguration())
		return;

	try
	{
		SetInitialisationFatalError("");

		LPCSTR names[] = {"OpenFile", "CreateFileA", "CreateFileW", "CloseHandle", "WriteFile", "WriteFileEx", "ReadFile", "ReadFileEx"};
		LPVOID hooks[] = {OpenFile_Hook, CreateFileA_Hook, CreateFileW_Hook, CloseHandle_Hook, WriteFile_Hook, WriteFileEx_Hook, ReadFile_Hook, ReadFileEx_Hook};
		LPVOID originalAdresses[8];

		Hook("Kernel32.dll", 8, names, hooks, originalAdresses);

		g_lpOpenFile = (LPOpenFile)originalAdresses[0];
		g_lpCreateFileA = (LPCreateFileA)originalAdresses[1];
		g_lpCreateFileW = (LPCreateFileW)originalAdresses[2];
		g_lpCloseHandle = (LPCloseHandle)originalAdresses[3];
		g_lpWriteFile = (LPWriteFile)originalAdresses[4];
		g_lpWriteFileEx = (LPWriteFileEx)originalAdresses[5];
		g_lpReadFile = (LPReadFile)originalAdresses[6];
		g_lpReadFileEx = (LPReadFileEx)originalAdresses[7];
	}
	catch (...)
	{
		Log("Exception thrown while initialising hooking dll");
	}
}

void UninitialiseHooks()
{
	try
	{
		//g_lpCloseHandle(g_hLogFile);
	}
	catch (...)
	{
		Log("Exception thrown while uninitialising hooking dll");
	}
}
